# encoding=utf-8
import threading
import cv2
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
from Enhancement.models.StillGAN.utils.util import tensor2im
from Enhancement.models.StillGAN.models.networks import get_norm_layer, ResUNet


def array2bytes(array_img, suffix):
    # 对数组的图片格式进行编码
    success, encoded_array = cv2.imencode("." + suffix, array_img)
    # 将数组转为bytes
    bytes_img = encoded_array.tostring()

    return bytes_img


class EnhanceWrapper:
    _inst = None
    _lock = threading.Lock()

    def __init__(self):
        self.transform = transforms.Compose([
            transforms.Resize((512, 512), Image.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.stillgan = ResUNet(3, 3, 64, norm_layer=get_norm_layer())
        self.model_parameter = torch.load("120_net_G_A.pth", map_location=self.device)
        self.stillgan.load_state_dict(self.model_parameter)
        self.stillgan = self.stillgan.to(self.device)
        self.stillgan.eval()

    def _preprocess(self, data_in: str):
        ori_Img = Image.open(data_in).convert("RGB")
        ori_size = ori_Img.size
        trans_image = self.transform(ori_Img)
        trans_dim = trans_image.unsqueeze(0)
        return trans_dim, ori_size

    @torch.no_grad()
    def infer(self, image_path: str):
        trans_dim, ori_size = self._preprocess(image_path)
        trans_cuda = trans_dim.to(self.device)
        res = self.infer_wrapper(trans_cuda, ori_size)
        return res

    def infer_wrapper(self, batch_data, ori_size):
        W, H = ori_size
        runned = self.stillgan(batch_data)
        image = tensor2im(runned)

        (r, g, b) = cv2.split(image)
        img_arr = cv2.merge([b, g, r])
        img_arr = cv2.resize(img_arr, (W, H), interpolation=cv2.INTER_CUBIC)
        img_bytes = array2bytes(img_arr, "png")

        image = np.asarray(bytearray(img_bytes), dtype=np.uint8)
        img_arr = cv2.imdecode(image, cv2.IMREAD_COLOR)
        return img_arr

    def __new__(cls, *args, **kwargs):
        EnhanceWrapper._lock.acquire()
        if cls._inst is None:
            cls._inst = super(EnhanceWrapper, cls).__new__(cls)
        EnhanceWrapper._lock.release()
        return cls._inst
